#!/bin/bash
# DeepGP 环境安装脚本 - 中国国内优化版
# 使用国内镜像源加速安装

echo "=========================================="
echo "  DeepGP 环境安装向导 (国内优化)"
echo "=========================================="
echo ""

# 获取脚本所在目录
SCRIPT_DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )"
cd "$SCRIPT_DIR"

# 环境名称
ENV_NAME="deepgp"
PYTHON_VERSION="3.11"

# 检查 conda
if ! command -v conda &> /dev/null; then
    echo "❌ 未检测到 Conda"
    echo "请先安装 Anaconda 或 Miniconda"
    echo "下载地址: https://mirrors.tuna.tsinghua.edu.cn/anaconda/miniconda/"
    exit 1
fi

echo "✓ 检测到 Conda"

# 配置 conda 国内镜像源
echo ""
echo "配置 Conda 镜像源 (清华大学)..."
conda config --add channels https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/free/
conda config --add channels https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main/
conda config --add channels https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/pytorch/
conda config --set show_channel_urls yes
echo "✓ Conda 镜像源配置完成"

# 配置 pip 国内镜像源
echo ""
echo "配置 pip 镜像源 (清华大学)..."
mkdir -p ~/.pip
cat > ~/.pip/pip.conf << EOF
[global]
index-url = https://pypi.tuna.tsinghua.edu.cn/simple
[install]
trusted-host = pypi.tuna.tsinghua.edu.cn
EOF
echo "✓ pip 镜像源配置完成"

# 检查环境是否已存在
if conda env list | grep -q "^${ENV_NAME} "; then
    echo ""
    echo "⚠️  环境 ${ENV_NAME} 已存在"
    read -p "是否删除并重新创建? (y/n) " -n 1 -r
    echo
    if [[ $REPLY =~ ^[Yy]$ ]]; then
        echo "正在删除旧环境..."
        conda env remove -n ${ENV_NAME} -y
    else
        echo "使用现有环境"
    fi
fi

# 创建环境
if ! conda env list | grep -q "^${ENV_NAME} "; then
    echo ""
    echo "步骤 1/5: 创建 Conda 环境 (Python ${PYTHON_VERSION})"
    echo "=========================================="
    conda create -n ${ENV_NAME} python=${PYTHON_VERSION} -y
    if [ $? -ne 0 ]; then
        echo "❌ 创建环境失败"
        exit 1
    fi
    echo "✓ 环境创建成功"
fi

# 激活环境
echo ""
echo "步骤 2/5: 激活环境"
echo "=========================================="
source "$(conda info --base)/etc/profile.d/conda.sh"
conda activate ${ENV_NAME}

if [ $? -ne 0 ]; then
    echo "❌ 激活环境失败"
    exit 1
fi
echo "✓ 环境已激活: ${ENV_NAME}"

# 检测 GPU
echo ""
echo "步骤 3/5: 检测 GPU"
echo "=========================================="
if command -v nvidia-smi &> /dev/null; then
    echo "✓ 检测到 NVIDIA GPU"
    nvidia-smi --query-gpu=name --format=csv,noheader | head -n 1
    
    # 获取 CUDA 版本
    CUDA_VERSION=$(nvidia-smi | grep "CUDA Version" | awk '{print $9}' | cut -d. -f1,2)
    echo "CUDA 版本: ${CUDA_VERSION}"
    
    if [[ "$CUDA_VERSION" > "11.7" ]] || [[ "$CUDA_VERSION" == "11.7" ]]; then
        TORCH_VERSION="torch torchvision torchaudio"
        DGL_VERSION="dgl -f https://data.dgl.ai/wheels/cu117/repo.html"
        echo "将安装 PyTorch 2.x (CUDA 11.7+)"
    else
        TORCH_VERSION="torch==1.12.1"
        DGL_VERSION="dgl -f https://data.dgl.ai/wheels/repo.html"
        echo "将安装 PyTorch 1.12.1 (兼容模式)"
    fi
    USE_GPU=true
else
    echo "⚠️  未检测到 GPU，将使用 CPU 版本"
    TORCH_VERSION="torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu"
    DGL_VERSION="dgl -f https://data.dgl.ai/wheels/repo.html"
    USE_GPU=false
fi

# 安装 PyTorch
echo ""
echo "步骤 4/5: 安装 PyTorch 和 DGL"
echo "=========================================="

if [ "$USE_GPU" = true ]; then
    echo "从清华镜像源安装 PyTorch (GPU版本)..."
    pip install ${TORCH_VERSION} -i https://pypi.tuna.tsinghua.edu.cn/simple
    
    if [ $? -eq 0 ]; then
        echo "✓ PyTorch 安装成功"
        
        # 验证 GPU 可用性
        python << EOF
import torch
if torch.cuda.is_available():
    print(f"✓ GPU 可用: {torch.cuda.get_device_name(0)}")
else:
    print("⚠️  PyTorch 安装成功但 GPU 不可用")
EOF
    else
        echo "❌ PyTorch 安装失败"
        echo "尝试手动安装:"
        echo "  pip install torch torchvision torchaudio -i https://pypi.tuna.tsinghua.edu.cn/simple"
        exit 1
    fi
else
    echo "从清华镜像源安装 PyTorch (CPU版本)..."
    pip install ${TORCH_VERSION} -i https://pypi.tuna.tsinghua.edu.cn/simple
    
    if [ $? -eq 0 ]; then
        echo "✓ PyTorch (CPU) 安装成功"
    else
        echo "❌ PyTorch 安装失败"
        exit 1
    fi
fi

# 安装 DGL
echo ""
echo "安装 DGL..."
pip install ${DGL_VERSION}

if [ $? -eq 0 ]; then
    echo "✓ DGL 安装成功"
else
    echo "⚠️  DGL 安装失败，尝试备用方法..."
    pip install dgl -f https://data.dgl.ai/wheels/repo.html
fi

# 安装其他依赖
echo ""
echo "步骤 5/5: 安装其他依赖"
echo "=========================================="

if [ -f "requirements.txt" ]; then
    echo "从清华镜像源安装基础依赖..."
    pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple
    if [ $? -eq 0 ]; then
        echo "✓ 基础依赖安装成功"
    else
        echo "⚠️  部分依赖安装失败，但可以继续"
    fi
fi

if [ -f "streamlit_requirements.txt" ]; then
    echo "从清华镜像源安装 Streamlit 依赖..."
    pip install -r streamlit_requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple
    if [ $? -eq 0 ]; then
        echo "✓ Streamlit 依赖安装成功"
    else
        echo "⚠️  Streamlit 依赖安装失败"
    fi
fi

# 验证安装
echo ""
echo "=========================================="
echo "  验证安装"
echo "=========================================="

python << EOF
import sys

packages = {
    'torch': 'PyTorch',
    'dgl': 'DGL',
    'transformers': 'Transformers',
    'pandas': 'Pandas',
    'numpy': 'NumPy',
    'FastNLP': 'FastNLP',
    'pyteomics': 'Pyteomics',
    'streamlit': 'Streamlit',
    'sklearn': 'scikit-learn'
}

print("\n安装包检查:")
success = True
for package, name in packages.items():
    try:
        mod = __import__(package)
        version = getattr(mod, '__version__', 'unknown')
        print(f"✓ {name:20s} {version}")
    except ImportError:
        print(f"❌ {name:20s} 未安装")
        success = False

# 检查 GPU 可用性
try:
    import torch
    print(f"\nPython 版本: {sys.version.split()[0]}")
    print(f"PyTorch 版本: {torch.__version__}")
    if torch.cuda.is_available():
        print(f"✓ GPU 可用: {torch.cuda.get_device_name(0)}")
        print(f"  GPU 数量: {torch.cuda.device_count()}")
        print(f"  CUDA 版本: {torch.version.cuda}")
    else:
        print("⚠️  使用 CPU 模式")
except:
    pass

if not success:
    print("\n⚠️  部分依赖安装失败")
    sys.exit(1)
EOF

if [ $? -eq 0 ]; then
    echo ""
    echo "=========================================="
    echo "  ✓ 安装完成！"
    echo "=========================================="
    echo ""
    echo "环境信息:"
    echo "  环境名称: ${ENV_NAME}"
    echo "  Python 版本: ${PYTHON_VERSION}"
    if [ "$USE_GPU" = true ]; then
        echo "  模式: GPU 加速"
    else
        echo "  模式: CPU"
    fi
    echo ""
    echo "下一步:"
    echo "  1. 激活环境: conda activate ${ENV_NAME}"
    echo "  2. 启动应用: ./start_app.sh"
    echo "  或直接运行: ./start_app.sh (会自动激活环境)"
    echo ""
    echo "命令行使用:"
    echo "  python 1_dataset_format.py --help"
    echo ""
    echo "文档:"
    echo "  查看 docs/快速开始.md"
    echo ""
else
    echo ""
    echo "=========================================="
    echo "  ⚠️  安装过程中出现问题"
    echo "=========================================="
    echo ""
    echo "请尝试手动安装:"
    echo "  conda activate ${ENV_NAME}"
    echo "  pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple"
    echo ""
fi

